import json
import os

from utils.draw_acc_ppl import moving_average_filter, _read_acc, draw_acc, draw_ppl

dataset_dir = "../../log"
interval = 2

datas = {
    "Transformer": None,
    "GAT+Transformer": None
}

dir_paths = []
for dir_path, dir_names, filenames in os.walk(dataset_dir):
    for dir_name in dir_names:
        if dir_name.startswith("log"):
            dir_paths.append(os.path.join(dataset_dir, dir_name))
dir_paths = dir_paths[:]

for i, dir_path in enumerate(dir_paths):
    acc_file = os.path.join(dir_path, "test_acc.txt")
    config_file = os.path.join(dir_path, "config.txt")
    if i % interval == 0:
        datas["GAT+Transformer"] = moving_average_filter(_read_acc(acc_file), 1)
    elif i % interval == 1:
        datas["Transformer"] = moving_average_filter(_read_acc(acc_file), 1)

        with open(config_file, 'r') as f:
            config = json.load(f)
        save_file_name = "acc_%d.svg" % (i // interval + 1)
        save_file_name = os.path.join(".", save_file_name)
        title = "seed=%d" % config["seed"]
        # title = ""
        draw_acc(datas, save_file=save_file_name, title=title)
        print("save file:", save_file_name)

for i, dir_path in enumerate(dir_paths):
    ppl_file = os.path.join(dir_path, "perplexity.txt")
    config_file = os.path.join(dir_path, "config.txt")
    if i % interval == 0:
        datas["GAT+Transformer"] = moving_average_filter(_read_acc(ppl_file), 1)
    elif i % interval == 1:
        datas["Transformer"] = moving_average_filter(_read_acc(ppl_file), 1)

        with open(config_file, 'r') as f:
            config = json.load(f)
        save_file_name = "ppl_%d.svg" % (i // interval + 1)
        save_file_name = os.path.join(".", save_file_name)
        title = "seed=%d" % config["seed"]
        # title = ""
        draw_ppl(datas, save_file=save_file_name, title=title)
        print("save file:", save_file_name)